import os
import cv2
import numpy as np
from scipy.stats import pearsonr

def calculate_pearson_correlation(img1, img2):
    img1 = img1.flatten()
    img2 = img2.flatten()
    return pearsonr(img1, img2)[0]

def process_folders(folder1, folder2, min_threshold, max_threshold):
    all_correlations = []
    filtered_correlations = []

    files1 = [f for f in os.listdir(folder1) if f.endswith('.png')]
    files2 = [f for f in os.listdir(folder2) if f.endswith('.png')]

    # Create a dictionary to map trx-try to file paths
    file_map1 = {f.split('_')[-1]: os.path.join(folder1, f) for f in files1}
    file_map2 = {f.split('_')[-1]: os.path.join(folder2, f) for f in files2}

    common_keys = set(file_map1.keys()).intersection(file_map2.keys())

    for key in common_keys:
        img1_path = file_map1[key]
        img2_path = file_map2[key]

        img1 = cv2.imread(img1_path, cv2.IMREAD_GRAYSCALE)
        img2 = cv2.imread(img2_path, cv2.IMREAD_GRAYSCALE)

        if img1 is not None and img2 is not None:
            pearson_corr = calculate_pearson_correlation(img1, img2)
            all_correlations.append(f"{key}:{pearson_corr}")

            if min_threshold <= pearson_corr <= max_threshold:
                filtered_correlations.append(f"{key}:{pearson_corr}")

    with open('all_correlations2.txt', 'w') as f:
        for item in all_correlations:
            f.write("%s\n" % item)

    with open('filtered_correlations2.txt', 'w') as f:
        for item in filtered_correlations:
            f.write("%s\n" % item)


# 设置文件夹路径和阈值
folder1 = 'E:\\wafer52\\11866_32nm'
folder2 = 'E:\\wafer52\\11867_32nm_fine_align'
min_threshold = 0.0
max_threshold = 1


process_folders(folder1, folder2, min_threshold, max_threshold)
